# -*- coding:utf-8 -*-
"""
双色球数据处理与模型训练/预测主程序
串联流程: 数据爬取 → 格式转换 → 特征工程 → 多模型训练/预测
支持通过--mode参数控制训练/预测模式
"""

import traceback
import warnings
import argparse  # 新增: 命令行参数解析
from config.ssq_config import BASE_URL, DATA_FOLDER, logger, SSQ_CONFIG
from trainer.spider_data import SSQLotterySpider
from trainer.format01_data import Lottery01Format
from trainer.merge_features import MergeFeatures
from trainer.multi_models_train import SSQMultiTrainer
from trainer.stacking_ensemble import StackingEnsemble
from trainer.stacking_BayesianOptimization import StackingBayesianOptimization
from predictor.predictor import SSQSingleColFusionPredictor
# 新增: 导入预测相关模块(假设已实现)
import trainer.utils
import threading

# 抑制多进程信号量泄露的警告
warnings.filterwarnings('ignore', category=UserWarning, module='multiprocessing.resource_tracker')


class ModelTrainThread(threading.Thread):
    """模型训练线程: 单独承载一个模型的完整训练流程(内部含多进程训练标签)"""
    def __init__(self, model_name: str, train_func):
        super().__init__()
        self.model_name = model_name  # 模型名称(如XGBoost)
        self.train_func = train_func  # 模型训练函数(如train_xgb)
        self.success = False  # 记录线程训练结果

    def run(self) -> None:
        """线程执行逻辑: 调用模型训练函数, 独立执行训练"""
        logger.info(f"[{self.model_name}训练线程]启动(线程ID: {threading.get_ident()})")
        try:
            # 执行模型训练(内部自动触发多进程训练33个标签)
            self.success = self.train_func()
            if self.success:
                logger.info(f"[{self.model_name}训练线程]训练完成! ")
            else:
                logger.error(f"[{self.model_name}训练线程]训练失败! ")
        except Exception as e:
            logger.error(
                f"[{self.model_name}训练线程]崩溃: {str(e)}\n{traceback.format_exc()}"
            )
            self.success = False


def crawl_ssq_data(run_type: str =  None):
    """爬取双色球原始数据"""
    logger.info("===== 1. 开始爬取双色球原始数据 =====")
    try:
        ssq_spider = SSQLotterySpider(run_type=run_type)
        if ssq_spider.run():
            logger.info(f"原始数据爬取完成.")
            return True  # 明确返回成功
        logger.info(f"原始数据爬取失败, 请检查日志!!")
        return False  # 明确返回成功
    except Exception as e:
        logger.critical(f"爬取过程异常:{str(e)},\n {traceback.format_exc()}" )
        return False  # 发生异常时返回失败

def format01_data(run_type: str =  None):
    """格式化数据为01序列"""
    logger.info("===== 2. 开始格式化数据为01序列 =====")
    try:
        # 初始化并运行01格式化工具
        format_tool = Lottery01Format(run_type=run_type)
        format_tool.run()

        if format_tool.run():
            logger.info(f"01序列化数据成功!")
            return True  # 明确返回成功
        logger.info(f"01序列化数据成功失败, 请检查日志!!")
        return False  # 明确返回成功
    except Exception as e:
        logger.critical(f"01序列化运行失败:{str(e)},\n {traceback.format_exc()}")
        return False

def feature_enginner(run_type: str =  None):
    try:
        # 集成分形特征和其他特征
        feature_engine = MergeFeatures(run_type=run_type)
        # 执行特征工程核心流程, 检查返回结果
        if not feature_engine.do_features():
            logger.critical("分形特征、30+特征生成及合并流程执行失败, 立即退出程序")
            return False
        
        logger.info("生成分形特征、30+特征、合并01数据及标签数据成功...")
        return True
    except Exception as e:
        logger.critical(
            f"生成分形特征、30+特征、合并01数据及标签数据异常: {str(e)},\n"
            f"异常堆栈: {traceback.format_exc()}"
        )
        return False

# 各模型独立训练入口
def train_xgb(run_type: str = None) -> bool:
    return SSQMultiTrainer(model_type="xgb").run(run_type=run_type)

def train_lgb(run_type: str = None) -> bool:
    return SSQMultiTrainer(model_type="lgb").run(run_type=run_type)

def train_lstm(run_type: str = None) -> bool:
    return SSQMultiTrainer(model_type="lstm").run(run_type=run_type)

def train_lr(run_type: str = None) -> bool:
    return SSQMultiTrainer(model_type="lr").run(run_type=run_type)

def train_rf(run_type: str = None) -> bool:
    return SSQMultiTrainer(model_type="rf").run(run_type=run_type)


def mutil_train(run_type: str =  None):
    try:
        """模型训练主入口"""
        logger.info("===== 多模型训练系统独立启动 =====")
                
        # 可选择需要运行的模型
        model_steps = [
            ("xgb", train_xgb),
            ("lgb", train_lgb),
            ("lstm", train_lstm),
            ("lr", train_lr),
            ("rf", train_rf),
        ]

        # 串行执行每个模型(内部已多进程训练标签)
        logger.info(f'1. 串行启动模型训练(每个模型内部多进程训练标签)')
        all_models_success = True
        failed_models = []
        for model_name, train_func in model_steps:
            logger.info(f"===== [{model_name.upper()}]模型训练start =====")
            if train_func(run_type=run_type):
                logger.info(f"===== [{model_name.upper()}]模型训练完成 =====")
            else:
                logger.error(f"===== {model_name.upper()}模型训练失败 =====")
                all_models_success = False
                failed_models.append(model_name)
        
        if not all_models_success:
            logger.critical(f"===== 训练失败的模型: {failed_models} =====")
            return False
        
        logger.info("===== 所有模型训练完成 =====")
        return True
    except Exception as e:
        logger.critical(f"模型训练异常终止: {str(e)},\n {traceback.format_exc()}")
        return False


def stacking_ensemble(run_type: str =  None):
    try:
        stacking = StackingEnsemble(run_type=run_type)
        summary = stacking.train()
        print(f"融合模型训练完成, 成功率: {summary['success_rate']:.2f}%")
        return True
    except Exception as e:
        logger.critical(f"融合模型训练异常终止: {str(e)},\n {traceback.format_exc()}")
        return False
    
def stacking_BayesianOptimization(run_type: str =  None):
    try:
        stacking = StackingBayesianOptimization(run_type=run_type)
        summary = stacking.train()
        print(f"融合模型训练完成, 成功率: {summary['success_rate']:.2f}%")
        return True
    except Exception as e:
        logger.critical(f"融合模型训练异常终止: {str(e)},\n {traceback.format_exc()}")
        return False


def run_prediction(run_type: str =  None):
    """预测流程: 复用特征工程, 加载模型输出结果"""
    logger.info(f"===== 进入预测模式 run_type={run_type}=====")
    try:       
        # 4. 加载模型并预测
        predictor = SSQSingleColFusionPredictor()
        if not predictor.run_predict():
            logger.info(f"预测失败，请检查! {traceback.format_exc()}")
            return False
        logger.info(f"预测已完成.")
        return True
    
    except Exception as e:
        logger.critical(f"预测流程异常终止: {str(e)}\n{traceback.format_exc()}")
        return False


def main():
    """主流程控制: 通过命令行参数--mode指定训练/预测模式"""
    # 新增: 解析命令行参数
    parser = argparse.ArgumentParser(description="双色球模型训练与预测系统")
    parser.add_argument(
        "--mode", 
        required=True, 
        choices=["train", "predict"],
        help="运行模式: train(训练模型)或 predict(预测下一期)"
    )
    args = parser.parse_args()

    logger.info(f"===== 双色球数据分析与模型{'训练' if args.mode == 'train' else '预测'}系统启动 =====")
    
    # 环境检查(训练和预测通用)
    if not trainer.utils.check_environment_with_gpu_detail():
        logger.critical("环境检查不通过, 程序终止")        
        return
    
    logger.info(f'BASE_URL: {BASE_URL}')
    logger.info(f'DATA_FOLDER: {DATA_FOLDER}')
    run_type = args.mode

    logger.info(f'main() run_type: {run_type}')

    steps = [
        # ("1.数据爬取", crawl_ssq_data),
        ("2.数据格式化", format01_data),
        ("3.30+维度特征工程", feature_enginner),   
    ]
    # 根据模式选择执行流程
    if args.mode == "train":
        # 训练模式流程
        steps += [
            ("4.模型训练", mutil_train),
            # ("5.模型融合", stacking_ensemble),
            ("5.模型融合", stacking_BayesianOptimization),
        ]
    else:
        # 预测模式流程(复用训练的特征工程, 新增预测专属步骤)
        steps += [
            ("4.预测流程", run_prediction)  # 预测流程已封装内部步骤
        ]

    # 执行对应模式的流程
    for step_name, step_func in steps:
        if not step_func(run_type=run_type):
            logger.critical(f"{step_name}失败, 程序终止,{traceback.format_exc()}")
            return
        logger.info(f"{step_name}已完成")
    
    logger.info(f"===== 所有{'训练' if args.mode == 'train' else '预测'}流程执行完成 =====")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logger.critical(f"主程序异常终止:{str(e)},{traceback.format_exc()}")
        exit(1)